// Copyright (c) 2009-2010, Jeremy Brewer
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice,
// this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
// LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
// CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
// SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
// INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
// CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
// ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
// POSSIBILITY OF SUCH DAMAGE.

package galaxie500.datastructures;

import java.util.ArrayList;
import java.util.Collections;

import galaxie500.datastructures.BoundingBox;
import galaxie500.distancemetrics.DistanceMetric;
import galaxie500.distancemetrics.EuclideanDistance;

/**
 * Pointerless KDTree implementation
 * 
 * KD-Trees are a data structure for efficiently performing localized searches
 * in multidimensional datasets. They are a natural extension of binary search
 * trees.
 * 
 * This tree is "pointerless" because it stores the nodes in an array rather
 * than using pointers within each node.
 * 
 * Thanks to Brent Bryan for his frequent and insightful discussions regarding
 * KDTree implementations.
 * 
 * @author Jeremy Brewer
 */
public class KDTree {

  /**
   * Class representing nodes in a KDTree.
   */
  public static class Node {
    public static final int NO_PIVOT = -1;

    public int minIndex = 0;
    public int maxIndex = 0;
    public BoundingBox boundingBox = null;
    public int pivotDimension = NO_PIVOT;

    public boolean isLeafNode() {
      return pivotDimension == NO_PIVOT;
    }
  }

  /** The dataset of points, a matrix-like object */
  private DataSet dataset;

  /** The dimension of each point */
  private int dimension = 0;

  /** The maximum number of points stored in a leaf node */
  private int maxLeafNodeSize = 20;

  /** The number of nodes created */
  private int numNodes = 0;

  /**
   * Indirect array of indices. We sort points in dataset by swapping indices
   * indirectly.
   */
  private int[] indices;

  /** Array of tree nodes */
  private KDTree.Node[] nodeBinaryTree;

  /** Distance metric used to evaluate distances between points */
  private DistanceMetric metric;

  /**
   * Creates a KDTree from the given dataset and with a Euclidean distance
   * metric. All points in dataset must have the same dimension.
   * 
   * @param dataset The dataset containing multidimensional points
   */
  public KDTree(DataSet dataset) {
    this(dataset, new EuclideanDistance());
  }

  /**
   * Creates a KDTree from the given dataset using the given distance metric.
   * All points in dataset must have the same dimension.
   * 
   * @param dataset The dataset containing multidimensional points
   * @param metric The distance metric to use when comparing points
   */
  public KDTree(DataSet dataset, DistanceMetric metric) {
    assert dataset.numPoints() > 0;
    this.dataset = dataset;
    this.dimension = dataset.getPoint(0).dimension();
    this.metric = metric;
    buildTree();
  }

  /**
   * Builds the tree.
   */
  private void buildTree() {
    // Sorting is done indirectly using this array.
    indices = new int[dataset.numPoints()];
    for (int i = 0; i < dataset.numPoints(); i++) {
      indices[i] = i;
    }

    int nodeBinaryTreeSize =
        computeMaxNodeIndex(dataset.numPoints(), maxLeafNodeSize);

    // The binary tree of nodes is represented as an array. When accessing this
    // array, we want getNode(1) to return index 0 of this array so that the
    // root node's children can be computed without a special case.
    nodeBinaryTree = new KDTree.Node[nodeBinaryTreeSize];

    numNodes = 0;
    buildNode(0, dataset.numPoints(), rootIndex());
  }

  /**
   * Given the total number of points and desired max leaf node size, computes
   * the maximum number of nodes needed. This is only possible because we always
   * create a balanced tree due to splitting on median values.
   * 
   * @param numPoints The number of points
   * @param maxLeafNodeSize The maximum number of points in a leaf node
   * @return The maximum require node index, i.e. the required number of nodes
   */
  private static int computeMaxNodeIndex(int numPoints, int maxLeafNodeSize) {
    // Start the loop 1 level below the root index.
    int numNonLeafNodes = 1;
    int numPointsInNode = numPoints / 2;
    if ((numPointsInNode % 2) != 0) numPointsInNode++;
    int numNodesOnThisLevel = 1;

    // Count the number of nodes. This relies on the fact that our KDTree
    // is always balanced because we pivot about the median.
    while (true) {
      numNonLeafNodes += 2 * numNodesOnThisLevel;
      numPointsInNode /= 2;
      numNodesOnThisLevel *= 2;
      if (numPointsInNode <= maxLeafNodeSize) break;
    }

    // The right child of the index for the rightmost non-leaf node gives
    // the max required node index (recall array is 1 offset).
    return rightChildIndex(numNonLeafNodes);
  }

  /**
   * Class for holding the results of a KDTree search. It is a simple struct
   * which contains the index of the point, as referenced in 'dataset' passed to
   * the KDTree constructor, and the Euclidean distance to that point.
   */
  public static class Match implements Comparable<Match> {
    public int index = -1;
    public double distance = Double.MAX_VALUE;

    public Match() {
    }

    public Match(int index, double distance) {
      this.index = index;
      this.distance = distance;
    }

    public int compareTo(Match that) {
      if (this.distance < that.distance) {
        return -1;
      } else if (this.distance > that.distance) {
        return 1;
      }
      return 0;
    }
  }

  /**
   * Returns all points which are within the given rectangular bounding box.
   * Upon return, 'closestPoints' will contain a sorted list of <index,
   * distance> giving the index of the points and its distance (as defined by
   * the distanceTo() method of 'point'), respectively.
   * 
   * @param point The point to search for
   * @param box The bounding box that this point must lie within
   * @param closestPoints The list of (index, distance) for the closest points
   */
  public void search(Point point, BoundingBox box,
                     ArrayList<KDTree.Match> closestPoints) {
    closestPoints.clear();
    closestPoints.ensureCapacity(50); // Avoid reallocs in inner loop.

    BoundingBox rootBoundingBox = getNode(rootIndex()).boundingBox;
    if (!box.intersects(rootBoundingBox)) return;
    searchImpl(point, box, rootIndex(), closestPoints);
    Collections.sort(closestPoints);
  }

  /**
   * Internal recursive searching routine.
   * 
   * @param point The point to search for
   * @param box The bounding box that this point must lie within
   * @param nodeIndex The index of the node to search
   * @param closestPoints The list of (index, distance) for the closest points
   */
  private void searchImpl(Point point, BoundingBox box, int nodeIndex,
                          ArrayList<KDTree.Match> closestPoints) {
    KDTree.Node node = getNode(nodeIndex);
    if (node.isLeafNode()) {
      // Recursive base case -- we are in a leaf node, so compute distances to
      // all points within this node.
      for (int i = node.minIndex; i < node.maxIndex; i++) {
        Point nodePoint = getPoint(i);
        if (box.contains(nodePoint)) {
          double distance = metric.distance(point, nodePoint);
          closestPoints.add(new Match(indices[i], distance));
        }
      }
    } else {
      // Check whether box intersects the children of this node.
      int leftIndex = leftChildIndex(nodeIndex);
      KDTree.Node leftChild = getNode(leftIndex);
      if (box.intersects(leftChild.boundingBox)) {
        searchImpl(point, box, leftIndex, closestPoints);
      }

      int rightIndex = rightChildIndex(nodeIndex);
      KDTree.Node rightChild = getNode(rightIndex);
      if (box.intersects(rightChild.boundingBox)) {
        searchImpl(point, box, rightIndex, closestPoints);
      }
    }
  }

  /**
   * Returns the index of the root node in nodeBinaryTree.
   * 
   * @return The root index
   */
  protected static int rootIndex() {
    return 1;
  }

  /**
   * Returns the left child of node i.
   * 
   * @param i The node index
   * @return The index of i's left child
   */
  protected static int leftChildIndex(int i) {
    return 2 * i;
  }

  /**
   * Returns the right child of node i.
   * 
   * @param i The node index
   * @return The index of i's right child
   */
  protected static int rightChildIndex(int i) {
    return 2 * i + 1;
  }

  /**
   * Returns the parent of node i.
   * 
   * @param i The node index
   * @return The index of i's parent
   */
  protected static int parentIndex(int i) {
    return i / 2;
  }

  /**
   * Returns the sibling of node i.
   * 
   * @param i The node index
   * @return The index of i's sibling
   */
  protected static int siblingIndex(int i) {
    if ((i % 2) == 0) return i + 1;
    return i - 1;
  }

  /**
   * Returns the ith point. We sort indirectly, so this helps with the index
   * gymnastics. Access to the underlying dataset should always be done using
   * this method.
   * 
   * @param i The ith point in the dataset to return.
   * @return The ith point
   */
  protected Point getPoint(int i) {
    return dataset.getPoint(indices[i]);
  }

  /**
   * Returns the ith node in nodeBinaryTree, allocating it first if necessary.
   * The underlying array is treated as 1 offset, i.e. having valid indices
   * ranging from 1 to nodeBinaryTree.length; this is done to avoid having
   * special case logic for the root node's children.
   * 
   * @param i The node index
   * @return The ith node
   */
  protected KDTree.Node getNode(int i) {
    int j = i - 1;
    if (nodeBinaryTree[j] == null) {
      nodeBinaryTree[j] = new KDTree.Node();
    }
    return nodeBinaryTree[j];
  }

  /**
   * @return The max leaf node size
   */
  public int getMaxLeafNodeSize() {
    return maxLeafNodeSize;
  }

  /**
   * @param maxLeafNodeSize The value to set the max leaf node size to
   */
  public void setMaxLeafNodeSize(int maxLeafNodeSize) {
    assert maxLeafNodeSize > 0 : "Invalid max leaf node size";
    this.maxLeafNodeSize = maxLeafNodeSize;
  }

  /**
   * Sorts the underlying dataset_ container indirectly over the range
   * [min_index, max_index) in 'indices' by swapping points based on their value
   * in dimension 'pivotDimension'. After the sort, it should be true that
   * dataset.getPoint(indices[i]).getValue(pivot_dimension) < pivot for all
   * indices i up to the median = (maxIndex + minIndex) / 2.
   * 
   * This method is taken from the select() function from "Algorithms in C++" by
   * Robert Sedgewick.
   * 
   * @param pivotDimension The dimension to partially sort about
   * @param minIndex The minimum index of points in 'dataset'
   * @param maxIndex The maximum index of points in 'dataset'
   */
  protected void partialSortAboutMedian(int pivotDimension, int minIndex,
                                        int maxIndex) {
    int left = minIndex;
    int right = maxIndex - 1;
    int median = (maxIndex + minIndex) / 2;

    while (right > left) {
      int i = left;
      int j = right - 1;
      double value = getPoint(right).getValue(pivotDimension);

      while (true) {
        // Find the first pair of points which aren't sorted.
        boolean foundSwap = false;

        while (i < right) {
          if (getPoint(i).getValue(pivotDimension) > value) {
            foundSwap = true;
            break;
          }
          i++;
        }

        while (j > left) {
          if (getPoint(j).getValue(pivotDimension) < value) {
            foundSwap = true;
            break;
          }
          j--;
        }

        if (i >= j || !foundSwap) break;
        swapPoints(i, j);
      }

      swapPoints(i, right);
      if (i >= median) right = i - 1;
      if (i <= median) left = i + 1;
    }
  }

  /**
   * Returns whether node i is valid.
   * 
   * @param i The node index
   * @return Whether the node index is valid
   */
  private boolean isValidNodeIndex(int i) {
    return i > 0 && i < nodeBinaryTree.length;
  }

  /**
   * Main recursive method which builds a node and then its children.
   * 
   * @param minIndex The minimum index of points in dataset
   * @param maxIndex The maximum index of points in dataset
   * @param nodeIndex The index of this node
   */
  protected void buildNode(int minIndex, int maxIndex, int nodeIndex) {
    numNodes++;
    assert isValidNodeIndex(nodeIndex) : "Invalid nodeIndex: " + nodeIndex;

    // Compute the bounding box for these points.
    BoundingBox boundingBox = computeBoundingBox(minIndex, maxIndex);

    // Set the properties of this node.
    Node node = getNode(nodeIndex);
    node.minIndex = minIndex;
    node.maxIndex = maxIndex;
    node.boundingBox = boundingBox;
    node.pivotDimension = Node.NO_PIVOT; // Leaf nodes have this value.

    // Check for the recursive base case when too few points remain.
    int numPointsInNode = maxIndex - minIndex;
    if (numPointsInNode <= maxLeafNodeSize) {
      return;
    }

    // Choose the pivot dimension for this level of the tree.
    int pivotDimension = boundingBox.largestDimension();

    // Corner case for recursion base case -- this occurs if the same point is
    // repeated more than maxLeafNodeSize times. We mark this as a leaf node.
    if (pivotDimension == -1) return;
    node.pivotDimension = pivotDimension;

    // Sort indirectly so that the median value for the pivotDimension is
    // located at i = (minIndex + maxIndex) / 2 and all points with
    // index < i have point[pivot_dimension] < median value.
    partialSortAboutMedian(pivotDimension, minIndex, maxIndex);

    // Recurse into each half (we already checked the base case above).
    int median = (minIndex + maxIndex) / 2;
    buildNode(minIndex, median, leftChildIndex(nodeIndex));
    buildNode(median, maxIndex, rightChildIndex(nodeIndex));
  }

  /**
   * Computes the bounding box for a node consisting of points with indices
   * ranging over [minIndex, maxIndex).
   * 
   * @param minIndex The minimum point index
   * @param maxIndex The maximum point index
   * @return A new BoundingBox instance describing the area the included points
   *         occupy
   */
  private BoundingBox computeBoundingBox(int minIndex, int maxIndex) {
    assert minIndex < maxIndex;
    double[] minValues = new double[dimension];
    double[] maxValues = new double[dimension];

    // Set the initial max/min values from the first point.
    for (int i = 0; i < dimension; i++) {
      minValues[i] = getPoint(minIndex).getValue(i);
      maxValues[i] = minValues[i];
    }

    // Search over the other points for larger or smaller values.
    for (int i = minIndex + 1; i < maxIndex; i++) {
      Point point = getPoint(i);
      for (int j = 0; j < dimension; j++) {
        double value = point.getValue(j);
        if (value < minValues[j]) minValues[j] = value;
        if (value > maxValues[j]) maxValues[j] = value;
      }
    }

    return new BoundingBox(minValues, maxValues);
  }

  /**
   * Indirectly swaps two points in the dataset.
   * 
   * @param i The ith point as referenced indirectly in 'indices'
   * @param j The jth point as referenced indirectly in 'indices'
   */
  private void swapPoints(int i, int j) {
    int tmp = indices[i];
    indices[i] = indices[j];
    indices[j] = tmp;
  }
}
